class CVAE(VAE):
    def __init__(self, latent_dim, num_classes):
        super(CVAE, self).__init__(latent_dim)
        self.num_classes = num_classes
        self.label_embedding = layers.Embedding(num_classes, latent_dim)

    def encode(self, x, labels):
        z_mean, z_log_var = super().encode(x)
        label_embeddings = self.label_embedding(labels)
        z_mean += label_embeddings
        z_log_var += label_embeddings
        return z_mean, z_log_var

    def decode(self, z, labels):
        label_embeddings = self.label_embedding(labels)
        z += label_embeddings
        return super().decode(z)

# 训练CVAE
cvae = CVAE(latent_dim=2, num_classes=10)
cvae.compile(optimizer='adam', loss=vae_loss)
cvae.fit([x_train, y_train], epochs=30, batch_size=128)